﻿type Color = 
    Red | Black | DoubleBlack
    
type RBTree<'a when 'a:comparison> =
    | Node of Color * 'a * int * RBTree<'a> * RBTree<'a> 
    | Leaf of Color



module RBTree = 
    let empty = Leaf (Black)

    let private getSize leftTree rightTree=
            match leftTree, rightTree with
                | Node (_, _, leftSize, _, _), Node (_, _, rightSize, _, _)->
                    leftSize+rightSize
                | Node (_, _, leftSize, _, _), Leaf _->
                    leftSize
                | Leaf _, Node (_, _, rightSize, _, _)->
                    rightSize
                | Leaf _, Leaf _->
                    0

    let rec contains x = function
        | Leaf _ -> false
        | Node (_, y, _, a, b) -> 
            if x = y then true
            elif x < y then contains x a
            else contains x b

    let rec depth = function 
        | Leaf _ -> 0
        | Node (_, _, _, a, b) -> 
            1 + max (depth a) (depth b)


    let insert x (t: RBTree<'a>) = 
        let balance = function
            | Black, z, size, Node (Red, y, _, Node (Red, x, _, a, b), c), d 
            | Black, z, size, Node (Red, x, _, a, Node (Red, y, _, b, c)), d
            | Black, x, size, a, Node (Red, z, _, Node (Red, y, _, b, c), d)
            | Black, x, size, a, Node (Red, y, _, b, Node (Red, z, _, c, d))->
                Node (Red, y, size, Node (Black, x, (getSize a b)+1, a, b), Node (Black, z, (getSize c d)+1, c, d))
            | color, b, _, c, d  -> // grandparent is red, does not change
                Node (color, b, (getSize c d)+1, c, d)

        let rec ins = function
            | Leaf _ ->
                Node (Red, x, 1, Leaf(Black), Leaf(Black))
            | Node (color, y, size, a, b) as s  ->
                if x < y then balance(color, y, (getSize a b)+2, ins a, b)
                elif x > y then balance(color, y, (getSize a b)+2, a, ins b)
                else s

        match ins t with
            | Node (Red, y, size, a, b)->
                Node (Black, y, size, a, b)
            | t -> t

    let rec min = function
        | Leaf _->
            None
        | Node (_, value, _, Leaf(Black), _)->
            Some(value)
        | Node (_, _, _, l, r)->
            min l

    let delete x t = 
        let isBlack = function   
            | Leaf(Black)
            | Node (Black, _, _, _, _)->
                true
            | _->
                false

        let isDoubleBlack = function
            | Leaf(DoubleBlack)
            | Node (DoubleBlack, _, _, _, _)->
                true
            | _-> 
                false

        let (++) color1 color2 =
            match color1,  color2 with
                | Red, Black
                | Black, Red ->
                    Black
                | Black, Black ->
                    DoubleBlack
                | _, _ -> 
                    Red

        let blackify = function
            | Leaf _ -> Leaf(Black)
            | Node (Black, _, _, _, _) as s -> s
            | Node (_, x, size, a, b) -> Node (Black, x, size, a, b)

        let rec balance = function
            // invariant: the two children rooted at node pv have same "black" depth
            // but the root node itself may become a "double" black node
            | Black, pv, _, x, Node (Red, rv, size, ra, rb) when (isDoubleBlack x)-> // case 1 
                balance(Black, rv, 0, balance(Red, pv, 0, x, ra), rb) // do left-rotate and continue balance
            | pc, pv, _, x, Node (Black, rv, size, ra, rb) when (isDoubleBlack x)-> // case 2.a 2.b and 2.c
                if isBlack ra && isBlack rb then // 2.a: reduce a black on both side
                    let tempNode=Node (Red, rv, size, ra, rb)
                    Node (pc++Black, pv, (getSize x tempNode)+1, blackify x, tempNode) // reduces the double node to root
                elif isBlack rb then // 2.b: do a right rotation in the right child, so rb becomes red and recudes to the "else" case
                    match ra with
                    | Node (_, rav, _, raa, rbb)->
                        let tempNode1= Node (Red, rv, (getSize rbb rb)+1, rbb, rb)
                        let tempNode2=Node (Black, rav, (getSize raa tempNode1)+1, raa, tempNode1)
                        balance(pc, pv, 0, x, tempNode2)
                    | _->
                        failwith "impossible error"
                else // 3.c 
                    let tempNode=Node (Black, pv, (getSize x ra)+1, blackify x, ra)
                    Node (pc, rv, (getSize tempNode rb)+1, tempNode, blackify rb)

            // when doubleblack x is on the right
            | Black, pv, _, Node (Red, lv, _, la, lb), x when (isDoubleBlack x)->
                balance(Black, lv, 0, la, balance(Red, pv, 0, lb, x))
            | pc, pv, _, Node (Black, lv, size, la, lb), x when (isDoubleBlack x)->
                if isBlack la && isBlack lb then
                    let tempNode=Node (Red, lv, (getSize la lb)+1, la, lb)
                    Node (pc++Black, pv, (getSize tempNode x)+1, tempNode, blackify x)
                elif isBlack la then
                    match lb with
                    | Node (_, _, lbv, lba, lbb)->
                        let tempNode1=Node (Red, lv, (getSize la lb)+1, la, lba)
                        let tempNode2=Node (Black, lbv, (getSize tempNode1 lbb)+1, tempNode1, lbb)
                        balance(pc, pv, 0, tempNode2, x)
                    | _->
                        failwith "impossible error"
                else   
                    let tempNode=Node (Black, pv, (getSize lb x)+1, lb, blackify x)
                    Node (pc, lv, (getSize la tempNode)+1, blackify la, tempNode)    
                 
            | a, b, _, c, d->
                Node (a, b, (getSize c d)+1, c, d)  


        let rec del value = function
            | Leaf(color) ->
                failwith "RBT delete failed because x is not in the tree"
            | Node (color, y, _, a, b) as s ->
                if value < y then balance(color, y, 0, del value a, b)
                elif value > y then balance(color, y, 0, a, del value b)
                else
                    match (a, b) with
                    | (Leaf(Black), Leaf(Black)) -> // without non-leaf children
                        Leaf(Black++color)
                    | (Leaf(Black), Node (nc, nv, size, nl, nr))
                    | (Node (nc, nv, size, nl, nr), Leaf(Black)) ->  // with a single child
                        Node (color++nc, nv, size, nl, nr)
                    | (l, r) ->
                        //find the successor (smallest element in a right child),
                        //replace the key with the successor's 
                        //and delete relevant node
                        match (min r) with
                        | Some(v) ->
                            balance(color, v, 0, l, del v r)
                        | None ->
                            failwith "impossible: can't find the successor"

        match del x t with
        | Node (_, y, size, a, b) ->
            Node (Black, y, size, a, b)
        | Leaf _ -> Leaf(Black)


    let ofList lst = 
        let rec ofList lst acc = 
            match lst with
                | [] -> acc
                | x::xs -> ofList xs (insert x acc)
        ofList lst empty

    let rec toList = function
        | Leaf _ -> []
        | Node (_, x, _, a, b) -> 
            (toList a) @ [x] @ (toList b)

    let rec map f = function
        | Leaf _ as l -> l
        | Node (color, x, size, l, r) -> Node (color, f x, size, map f l, map f r)

    let rec iter f = function
        | Leaf _ -> ()
        | Node (color, x, size, l, r) ->
            iter f l
            f x
            iter f r
        
    let rec fold f t acc = 
        match t with
        | Leaf _ -> acc
        | Node (color, x, size, l, r) -> fold f l (f x (fold f r acc))



    let size = function
        | Leaf _ -> 0
        | Node (_, _, size, _, _) -> size

    let nth n t = 
        let rec nth' n = function
            | Leaf _->
                failwith "can't find the nth member"
            | Node (_, v, size, l, r)->
                let lsize=getSize l (Leaf(Black))
                if lsize+1=n then
                    v
                elif lsize+1>n then
                    nth' n l
                else
                    nth' (n-lsize-1) r
        if n >= size t || n < 0 then
            failwith "nth out of range"
        else
            nth' (n+1) t

    let validate t = 
        let rec validate = function  
            | Leaf Black -> true, 0
            | Leaf DoubleBlack -> 
                printfn "DoubleBlack in a Leaf node"
                false, 0
            | Leaf Red ->
                printfn "Red in a Leaf node"
                false, 0
            | Node (Black, _, _, a, b) -> 
                let lr, ln = validate a
                let rr, rn = validate b
                lr && rr && (ln = rn), ln + 1
            | Node (Red, _, _, a, b) ->
                match a, b with
                    | Node (Red, _, _, _, _), _
                    | _, Node (Red, _, _, _, _) ->
                        let lr, ln = validate a
                        printfn "Red-Red in a tree"
                        false, ln
                    | _ -> 
                        let lr, ln = validate a
                        let rr, rn = validate b
                        lr && rr && (ln = rn), ln
            | Node (DoubleBlack, _, _, a, _) -> 
                let lr, ln = validate a
                printfn "DoubleBlack in an internal node"
                false, ln
        match t with
        | Leaf Black -> true, 0
        | Leaf _ -> false, 0
        | Node (Black, _, _, _, _) ->
            validate t
        | Node (Red, _, _, _, _)
        | Node (DoubleBlack, _, _, _, _) ->
            printfn "root must be Black"
            false, 0



module Test1 =         

    let t = RBTree.ofList [1..1000]

    for i=1 to 1000 do 
        printfn "%dth = %d" i (RBTree.nth i t)

    RBTree.nth 100 t

    let t1 = RBTree.delete 20 t
    let t2 = RBTree.delete 100 t1

    RBTree.toList t
    RBTree.toList t1
    RBTree.toList t2

    RBTree.nth 100 t
    RBTree.nth 100 t1
    RBTree.nth 100 t2
    RBTree.nth 99 t2
    RBTree.nth 98 t2
    RBTree.nth 2 t2

    RBTree.validate t
    RBTree.validate t1
    RBTree.validate t2

    RBTree.depth (RBTree.ofList [1..1000000])



#time

module Test2 = 
    // insert from 1 to n, and then delete from 1 to n
    let insertAndDelete n = 
        let rec insertall n rb i = 
            printfn "%A" (RBTree.validate rb)
            if i < n then
                insertall n (RBTree.insert i rb) (i+1)
            else
                rb

        let rec deleteall n rb i = 
            printfn "%A" (RBTree.validate rb)
            if i < n then
                deleteall n (RBTree.delete i rb) (i+1)
            else
                rb

        let rb = insertall n RBTree.empty 0
        let rb0 = deleteall n rb 0
        rb, rb0

    insertAndDelete 1000


module Test3 = 
    // Fisher-Yates shuffle
    let randperm n = 
        let p = Array.init n (fun i -> i)
        let r = new System.Random(1)
        for i=n-1 downto 1 do
            let j = r.Next(i+1) // random number 0 <= j <= i
            let tmp = p.[j]
            p.[j] <- p.[i]
            p.[i] <- tmp
        p

    let lst = randperm 1000000 |> Seq.toList

    let rb = RBTree.ofList lst |> ignore

    open Microsoft.FSharp.Collections

    let av = Set.ofList lst |> ignore

    //timing on a AMD E-350 Netbook (1.6 GHz)
    //    Real: 00:00:44.635, CPU: 00:00:44.694, GC gen0: 218, gen1: 61, gen2: 4
    //
    //    val rb : unit = ()
    //
    //    > 
    //    Real: 00:00:17.448, CPU: 00:00:18.766, GC gen0: 76, gen1: 20, gen2: 2
    //
    //    val av : unit = ()
    //







module Travel = 
    let rec map f = function
        | Leaf _ as l -> l
        | Node (color, x, size, l, r) -> Node (color, f x, size, map f l, map f r)

    // use explicit continuation
    let map2 f t = 
        let rec map f t k = 
            match t with
            | Leaf _ as l -> k l
            | Node (color, x, size, l, r) -> 
                map f l (fun lmap -> 
                    map f r (fun rmap -> k (Node(color, f x, size, lmap, rmap))))

        map f t (fun x -> x)

        

    type Cont<'a,'r> = ('a -> 'r) -> 'r

    type ContBuilder() =
        member x.Return (a):Cont<'a,'r> = fun k -> k a
        member x.Bind (c:Cont<'a,'r>, f:'a -> Cont<'b,_>) =
            (fun k -> c (fun a -> f a k))
        member this.Delay(f) = f()

    let cont = ContBuilder()

    let map3 f t = 
        let rec map f = function
            | Leaf _ as t -> cont { return t }
            | Node (color, x, size, l, r) -> 
                cont {
                    let! lm = map f l
                    let! rm = map f r
                    return Node (color, f x, size, lm, rm)
                }
        map f t (fun x -> x)
    

    let t = RBTree.ofList [1;2;3;4;5]

    let t1 = map (fun x -> x*x) t

    let t2 = map2 (fun x -> x*x) t

    let t3 = map3 (fun x -> x*x) t 



    
